package net.sansi.v3correctserver.service.correct.bright;

import cn.hutool.core.date.StopWatch;
import cn.hutool.json.JSONUtil;
import lombok.Data;
import lombok.Getter;
import lombok.experimental.Accessors;
import lombok.extern.slf4j.Slf4j;
import net.sansi.v3correctserver.entity.IntPoint;
import net.sansi.v3correctserver.entity.image.RecImage;
import net.sansi.v3correctserver.entity.photorec.LedNode;
import net.sansi.v3correctserver.exception.ServiceException;
import net.sansi.v3correctserver.service.correct.RectCorrectBase;
import net.sansi.v3correctserver.service.photorec.rec.RecRectScreen;
import net.sansi.v3correctserver.util.CorrectUtil;
import net.sansi.v3correctserver.util.Matrix;
import net.sansi.v3correctserver.util.MonotoneUtil;
import net.sansi.v3correctserver.util.RGBUtil;

import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * Created by Fang
 * create time : 2024/11/4
 * IDE : IntelliJ IDEA
 * <p>
 * 长方形屏幕亮度矫正
 */

@Slf4j
public class RectBrightCorrect extends RectCorrectBase {

    // 参数
    private Params params;

    public RectBrightCorrect(Params params) {
        this.params = params;
        initBaseParams(
                params.getResult().getWidth(),
                params.getResult().getHeight(),
                params.getSpan()
        );
    }

    // 高亮照片识别
    private RecRectScreen rec1;
    // 低量照片识别
    private RecRectScreen rec2;
    // 亮度矫正斜率
    private double[] k = {0, 0, 0};

    // 目标亮度
    @Getter
    private double[] targetPhotoBright = {0, 0, 0};
    // 颜色索引
    private int colorIdx = -1;

    public void toCorrect(RecImage highPhoto, RecImage lowPhoto, int colorIdx) {
        log.info("开始亮度矫正 [color] {} ----------------", RGBUtil.colorName(colorIdx));
        log.info("[showBrightnesses] {}", JSONUtil.toJsonStr(params.showBrightness));
        recSize = CorrectUtil.computeRecSize(showWidth, showHeight, span, offsetX, offsetY);
        StopWatch watch = new StopWatch("相片亮度矫正");
        watch.start("图像加载和识别");
        this.colorIdx = colorIdx;
        rec1 = recPhoto(highPhoto);
        rec2 = recPhoto(lowPhoto);
        if (rec1 == null) throw new ServiceException("需要高亮度相片");
        if (rec2 == null && k[colorIdx] == 0) throw new ServiceException("需要低亮度相片");
        watch.stop();
        log.info("[图像识别耗时] {}ms", watch.getLastTaskTimeMillis());
        watch.start("亮度矫正");
        k[colorIdx] = computeK();
        targetPhotoBright[colorIdx] = computeTargetPhotoBright();
        initKList();
        log.info("[K] {} [targetBrightness] {}", k, targetPhotoBright[colorIdx]);
        correctBright();
        watch.stop();
        log.info("[矫正耗时] {}ms", watch.getLastTaskTimeMillis());
        log.info("---------------- 完成亮度矫正 [{}ms]", watch.getTotal(TimeUnit.MILLISECONDS));
    }

    public void toCorrect(Path highPhoto, Path lowPhoto, int colorIdx) {
        RecImage high = RecImage.load(highPhoto);
        RecImage low = lowPhoto != null ? RecImage.load(lowPhoto) : null;
        toCorrect(high, low, colorIdx);
    }


    /**
     * 识别图像
     *
     * @param photo 图片
     * @return RecRectScreen
     */
    private RecRectScreen recPhoto(RecImage photo) {
        if (photo == null) return null;
        RecRectScreen rec = new RecRectScreen(photo, params.getBlackPhoto(), params.getFence()).toRec();
        if (rec.getMatrix().getWidth() != recSize.getWidth() || rec.getMatrix().getHeight() != recSize.getHeight()) {
            throw new ServiceException("识别到的行列与理论行列不一致！[偏移] {},{} [理论] {},{} [实际] {},{}",
                    offsetX, offsetY,
                    recSize.getWidth(),
                    recSize.getHeight(),
                    rec.getMatrix().getWidth(),
                    rec.getMatrix().getHeight()
            );
        }
        return rec;
    }

    /**
     * 执行亮度矫正
     */
    private void correctBright() {
        int showHighBrightness = params.showBrightness[colorIdx][0];
        rec1.getMatrix().foreach((mx, my, led1) -> {
            if (led1 == null) return;
            double k1 = findK(led1.getBrightness());
            double res = showHighBrightness - (led1.getBrightness() - targetPhotoBright[colorIdx]) * k1;
            double scale = res / showHighBrightness;
            scale = Math.pow(scale, params.attenuation);
            //scale = (scale - 1) / 2 + 1;
            //scale = Math.pow(Math.abs(scale), 0.9) * ((scale >= 1) ? 1 : -1);
            setResult(mx, my, scale);
            if (params.ledBorders != null) {
                IntPoint p = mxyToRxy(mx, my);
                params.ledBorders.set(p, led1.borderPixels);
            }
        });
    }

    /**
     * 设置结果
     *
     * @param mx    rec矩阵中 x 坐标
     * @param my    rec矩阵中 y 坐标
     * @param scale 矫正系数
     */
    private void setResult(int mx, int my, double scale) {
        IntPoint p = mxyToRxy(mx, my);
        double[] rgb = params.result.get(p);
        if (rgb == null) {
            rgb = new double[3];
            params.result.set(p, rgb);
        }
        rgb[colorIdx] = scale;
    }

    /**
     * 计算目标亮度
     */
    private double computeTargetPhotoBright() {
        log.info("[target photo bright] {}", params.getTargetPhotoBright());
        if (params.getTargetPhotoBright()[colorIdx] != -1) return params.getTargetPhotoBright()[colorIdx];
        //if (targetPhotoBrightness[colorIdx] != 0) return targetPhotoBrightness[colorIdx];
        int p = (int) Math.round(Math.max(rec1.getMatrix().getItemSize() * params.getBrightDecay(), 1));
        MonotoneUtil<Double> mono = new MonotoneUtil<>(1, Comparator.comparingDouble(a -> a));
        rec1.getMatrix().foreach((mx, my, led1) -> {
            if (led1 == null) return;
            mono.add(led1.getBrightness());
        });
        double mid = rec1.getLedByBrightSort(0.5).getBrightness();
        log.info("[photoBrightness] [mono] {} [mid] {} [decayed] {}", mono.get(), mid, mid * params.getBrightDecay());
        double tgBright = Math.max(mono.get(), mid * params.getBrightDecay());
        params.getTargetPhotoBright()[colorIdx] = tgBright;
        return tgBright;
        //return mono.get();
    }

    // 三色系数
    private List[] kLists = {
            new ArrayList<double[]>(),
            new ArrayList<double[]>(),
            new ArrayList<double[]>(),
    };

    private void initKList() {
        if (rec2 == null) return;
        List<double[]> list = new ArrayList<>();
        kLists[colorIdx] = list;
        int showHigh = params.showBrightness[colorIdx][0];
        int showLow = params.showBrightness[colorIdx][1];
        int showDiff = showHigh - showLow;
        int[] lessCount = new int[]{0};
        rec1.getMatrix().foreach((mx, my, led1) -> {
            if (led1 == null) return;
            LedNode led2 = rec2.getMatrix().get(mx, my);
            if (led2 == null) return;
            double b1 = led1.getBrightness();
            double b2 = led2.getBrightness();
            if (b1 <= b2) {
                lessCount[0]++;
                //return;
            }
            double k = showDiff / (b1 - b2);
            list.add(new double[]{b1, k});
        });
        log.info("[灯点总数] {} [逆亮度数量] {}", rec1.getMatrix().getItemSize(), lessCount);
        list.sort(Comparator.comparingDouble(a -> a[0]));
        smooth(1);
        for (double i = 0.1; i < 1; i += 0.1) {
            int idx = (int) Math.round((list.size() - 1) * i);
            double[] item = (double[]) kLists[colorIdx].get(idx);
            log.info("[bright] {} [k] {}", item[0], item[1]);
        }
    }

    /**
     * 平滑k系数
     *
     * @param count 递归次数
     */
    private void smooth(int count) {
        if (count < 1) return;
        int radius = 20;
        List<double[]> smoothList = new ArrayList<>();
        List<double[]> list = kLists[colorIdx];
        int n = kLists[colorIdx].size();
        for (int i = 0; i < n; i++) {
            double[] sc = {0, 0};
            for (int j = i - radius; j <= i + radius; j++) {
                if (j < 0 || j >= n) continue;
                sc[0] += list.get(j)[1];
                sc[1]++;
            }
            smoothList.add(new double[]{list.get(i)[0], sc[0] / sc[1]});
        }
        kLists[colorIdx] = smoothList;
        smooth(count - 1);
    }

    /**
     * 查找 k 系数
     *
     * @param brightness 亮度
     * @return k 系数
     */
    private double findK(double brightness) {
        List<double[]> list = kLists[colorIdx];
        int l = 0, r = list.size() - 1;
        while (l <= r) {
            int mid = (l + r) / 2;
            double[] p = list.get(mid);
            if (p[0] == brightness) return p[1];
            if (p[0] > brightness) r = mid - 1;
            else l = mid + 1;
        }
        return list.get(Math.max(r, 0))[1];
    }

    /**
     * 计算斜率 k 系数
     */
    private double computeK() {
        if (k[colorIdx] != 0 && rec2 == null) return k[colorIdx];
        int showHigh = params.showBrightness[colorIdx][0];
        int showLow = params.showBrightness[colorIdx][1];
        int showDiff = showHigh - showLow;
        double[] sum = {0, 0, 0};
        rec1.getMatrix().foreach((mx, my, led1) -> {
            if (led1 == null) return;
            LedNode led2 = rec2.getMatrix().get(mx, my);
            if (led2 == null) return;
            sum[0] += led1.getBrightness();
            sum[1] += led2.getBrightness();
            sum[2]++;
        });
        return showDiff / (sum[0] / sum[2] - sum[1] / sum[2]);
    }

    /**
     * 参数
     */
    @Data
    @Accessors(chain = true)
    public static class Params {
        private int span;
        // 衰减系数，对最终结果做一个指数操作
        private double attenuation = 1;
        // [colorIdx] = [high, low];
        private int[][] showBrightness;
        private RecImage blackPhoto;
        private List<IntPoint> fence;
        // [x,y] = 百分比 宽高为 showArea 宽高
        private Matrix<double[]> result;
        // [x, y] = [set of pixel]  宽高为 showArea 宽高
        private Matrix<Set<Long>> ledBorders;
        private double allowEmptyPercent = 0.5;
        // 基准亮度百分位
        private double brightDecay = 1.0;
        // 目标相片亮度
        private double[] targetPhotoBright = new double[]{-1, -1, -1};
    }

}
